import cv2
import numpy as np


def createFeatureDetector(name='SURF'):
    if name == 'SIFT':
        detector = cv2.xfeatures2d.SIFT_create()
    elif name == 'SURF':
        detector = cv2.xfeatures2d.SURF_create()
    return detector


def detect(gray):
    # gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    detector = createFeatureDetector()
    keypoints, descriptor = detector.detectAndCompute(gray, None)
    return keypoints, descriptor


def bf_knnmatches(matches,kp1, kp2):
    MIN_MATCH_COUNT = 10
    # store all the good matches as per Lowe's ratio test.
    good = []
    matchesMask = []
    goodsMatches = []
    if len(matches[0]) == 2:
        for i, (m, n) in enumerate(matches):
            if m.distance < 0.8 * n.distance:
                good.append(m)
                goodsMatches.append(matches[i])
        if len(good) > MIN_MATCH_COUNT:
            src_pts = np.float32([kp1[m.queryIdx].pt for m in good]).reshape(-1, 1, 2)
            dst_pts = np.float32([kp2[m.trainIdx].pt for m in good]).reshape(-1, 1, 2)
            M, mask = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC, 5.0)
            if M is not None:
                mask_list = mask.ravel().tolist()

                for i in range(len(mask_list)):
                    matchesMask.append([mask_list[i], 0])

    return goodsMatches, matchesMask


def detect_and_matches(img1, img2):
    img1 = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY)
    img2 = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY)

    k1, d1 = detect(img1)
    k2, d2 = detect(img2)

    matcher = cv2.BFMatcher()
    matches = matcher.knnMatch(d1, d2, k=2)
    goodMatches, matchesMask = bf_knnmatches(matches, k1, k2)

    draw_params = dict(matchColor=(0, 255, 0), singlePointColor=(255, 0, 0), matchesMask=matchesMask, flags=0)
    img3 = cv2.drawMatchesKnn(img1, k1, img2, k2, goodMatches, None, **draw_params)
    print(len(goodMatches))
    cv2.imshow('aaa', img3)
    cv2.waitKey()

    return goodMatches, matchesMask


if __name__ == '__main__':
    img1 = cv2.imread('/Volumes/bakup/cloud/matrix/image/blocks/02.png')
    img2 = cv2.imread('/Volumes/bakup/cloud/matrix/image/blocks/03.png')
    goodMatches, matchesMask = detect_and_matches(img1, img2)

    # draw_params = dict(matchColor=(0, 255, 0), singlePointColor=(255, 0, 0), matchesMask=matchesMask, flags=0)
    # img3 = cv2.drawMatchesKnn(img1, k1, img2, k2, goodMatches, None, **draw_params)
    # print(len(goodMatches))
    # cv2.imshow('aaa', img3)
    # cv2.waitKey()
